from typing import Dict, List
from .node_handler import NodeHandler,FollowHandler
from ..generator import ElementwiseGenerator, StrategyGenerator
from ..StrategiesVector import StrategiesVector
from .registry import operator_registry
from ..shard.placement_types import DeviceMesh,AbstractTensor,OperationData, OperationDataType
from geesibling.core.types import Graph, Node


__all__ = ['ElementwiseHandler']


@operator_registry.register('aten.gelu.default')
@operator_registry.register('aten.add.Tensor')
@operator_registry.register('aten.clone.default')
@operator_registry.register('aten.split_with_sizes.default')    #考虑到切完都是一样的切分策略
@operator_registry.register('aten.div.Tensor')
@operator_registry.register('aten.slice.Tensor')
@operator_registry.register('aten.slice_scatter.default')
@operator_registry.register('aten._softmax.default')
@operator_registry.register('aten.mul.Tensor')
@operator_registry.register('aten.pow.Tensor_Scalar')
@operator_registry.register('aten.tanh.default')
@operator_registry.register('aten.native_dropout.default')
@operator_registry.register('aten._to_copy.default')
@operator_registry.register('aten.copy.default')
@operator_registry.register('aten.sub.Tensor')
@operator_registry.register('aten.ge.Scalar')
@operator_registry.register('aten.eq.Scalar')

@operator_registry.register('aten.lt.Tensor')   # 2025.5.27 new add
@operator_registry.register('aten.le.Scalar')   # 2025.5.27 new add
@operator_registry.register('aten.logical_and.default')   # 2025.5.27 new add
@operator_registry.register('aten.full_like.default')       # 2025.5.27 new add

@operator_registry.register('aten.gt.Tensor')
@operator_registry.register('aten.mean.dim')
@operator_registry.register('aten.rsqrt.default')
@operator_registry.register('aten.cat.default') #可以换reshape
@operator_registry.register('aten.cos.default')
@operator_registry.register('aten.sin.default')
@operator_registry.register('aten.neg.default')
@operator_registry.register('aten.mul.Scalar')

@operator_registry.register('aten.select.int')
@operator_registry.register('aten.index.Tensor')

@operator_registry.register('aten._scaled_dot_product_efficient_attention.default') # 不确定
@operator_registry.register('aten.sigmoid.default')

@operator_registry.register('aten.argmax.default')
@operator_registry.register('aten.remainder.Scalar')


class ElementwiseHandler(FollowHandler):
    """
    A UnaryElementwiseHandler which deals with the sharding strategies for Elementwise Op.
    """

    def get_strategy_generator(self) -> List[StrategyGenerator]:
        op_data_mapping = self.get_operation_data_mapping()
        generators = []
        generators.append(ElementwiseGenerator(op_data_mapping, self.device_mesh, self.pre_strategies_vectors))
        return generators

    def get_operation_data_mapping(self) -> Dict[str, OperationData]:
        # 一个输入一个输出
        # 得到node的输入node
        mapping = {}
        for i in range(len(self.node.inputs)):
            if self.node.input_shape(i):
                input_abs_data = AbstractTensor(self.node.input_shape(i),self.node.input_type(i))
                physical_input = OperationData(name=str(self.node.inputs[i]), type=OperationDataType.ARG, data=input_abs_data)
                mapping["input_{}".format(i)] = physical_input
        output_abs_data = AbstractTensor(self.node.output_shape(0),self.node.output_type(0))
        physical_output = OperationData(name=str(self.node), data=output_abs_data,type=OperationDataType.OUTPUT)
        mapping['output'] = physical_output
        return mapping
